#include "pointcloud_tools.h"
#include "point_preprocess.h"
#include <iostream>
#include "point_common.h"
#include <fstream>

using namespace pointcloud_tools;
using namespace PointPreprocess;

int main(int argc, char **argv)
{
    Timer timer, timer_preprocess, timer_wt_init, timer_run;
    timer.reset();
    timer_preprocess.reset();
    timer_wt_init.reset();
    timer_run.reset();
    timer.start();

#ifdef SIMULATION
    Pointcloud_utils ptu;
    ptu.ddrSize = 8 * 8 * 1024 * 1024 * sizeof(int); // 256MB
    assert(ptu.Model_DDR_alloc("DDR_32MB.bin"));
    printf("Simulation Mode is ON!\n");
#else
    Pointcloud_online_utils ptu;
    ptu.pmem_mmap_init();
    printf("Onchip Mode is ON!\n");
#endif
    ptu.DDR_BASE_ADDR = 0x95000000;

    /////////////////////////////////////
    // layer1                          //
    // SUBCONV: ICH = 4, OCH = 32      //
    /////////////////////////////////////
    timer_preprocess.start();
    ptu.layer_config.cfg_inch_loop_num = 0;
    ptu.layer_config.cfg_core_enable = 0xf; // 32 output channels
    ptu.layer_config.cfg_shift_num = 8;
    ptu.layer_config.BASE_WT_ADDR = ptu.DDR_BASE_ADDR;
    size_t dlen = readBinToVec(ptu.layer_config.wt_data, "../data/golden_data/conv1_weight_packed_pad.bin");
    assert(dlen > 0);
    BlockConfig bcfg;
    dlen = readBinToVec(bcfg.pt_data, "../data/golden_data/voxel_coords.bin");
    assert(dlen > 0);
    dlen = readBinToVec(bcfg.in_data, "../data/golden_data/voxel_features.bin");
    assert(dlen > 0);

    uint16_t valid_pt_num = bcfg.pt_data.size();
    bcfg.cfg_cam_valid_num = valid_pt_num;
    bcfg.cfg_total_point_num = valid_pt_num;
    bcfg.BASE_PT_ADDR = ptu.DDR_BASE_ADDR + 0x01000000;
    bcfg.BASE_IN_ADDR = ptu.DDR_BASE_ADDR + 0x02000000;
    bcfg.BASE_RS_ADDR = ptu.DDR_BASE_ADDR + 0x03000000;

    uint16_t spatial_shape[6] = {0, 0, 0, 800, 600, 800}; // [z_min, y_min, x_min, z_max, y_max, x_max]

    std::vector<std::vector<int>> pt_mapping;
    std::vector<BlockConfig> bcfg_partition;
    std::vector<std::vector<int>> valid_filter;

    // Partition the point cloud into blocks
    octree_partition(bcfg, 1024,
                     pt_mapping, bcfg_partition, valid_filter,
                     ptu.layer_config.cfg_cam_search_range,
                     spatial_shape,
                     ptu.layer_config.cfg_inch_loop_num, 3);

    export_blocks_to_files(
        bcfg_partition, pt_mapping, valid_filter,
        "output_blocks",
        bcfg.pt_data, bcfg.in_data);

    for (int i = 0; i < bcfg_partition.size(); i++)
    {
        ptu.block_config.push_back(bcfg_partition[i]);
    }

    timer_preprocess.stop();
    timer_preprocess.print("Preprocess time");

    timer_wt_init.start();
    ptu.wt_init();
    timer_wt_init.stop();
    timer_wt_init.print("Weight init time for layer1");

    // Use a fixed block size of 1024 to get instructions and reuse them
    BlockConfig block_pt_1024;
    block_pt_1024.cfg_cam_valid_num = 1024;
    block_pt_1024.cfg_total_point_num = 1024;
    block_pt_1024.BASE_PT_ADDR = bcfg.BASE_PT_ADDR;
    block_pt_1024.BASE_IN_ADDR = bcfg.BASE_IN_ADDR;
    block_pt_1024.BASE_RS_ADDR = bcfg.BASE_RS_ADDR;
    int cnn_inst_index = ptu.inst_gen_block(block_pt_1024);

#ifdef SIMULATION
    ptu.compute_init();
    ptu.run();
#else
    timer_run.start();
    ptu.run_block(cnn_inst_index);
    timer_run.stop();
    timer_run.print("Run time for layer1");
#endif

    std::vector<uint64_t> layer1_merged_res;
    merge_RESULT_feat(ptu.block_config, bcfg.cfg_cam_valid_num, 3, pt_mapping, valid_filter, layer1_merged_res);

#ifdef SIMULATION
    std::string layer1_output_file = "layer1_output.txt";
    exportResults(layer1_merged_res, valid_pt_num, 3, layer1_output_file);
#endif

    //////////////////////////////////////
    // layer2                           //
    // SubConv: ich = 32, och = 8       //
    //////////////////////////////////////

    ptu.layer_config.cfg_inch_loop_num = 3;
    ptu.layer_config.cfg_core_enable = 0x1; // 8 output channels
    ptu.layer_config.cfg_shift_num = 8;
    ptu.layer_config.BASE_WT_ADDR = ptu.DDR_BASE_ADDR;
    dlen = readBinToVec(ptu.layer_config.wt_data, "../data/golden_data/conv2_weight_packed_pad.bin");
    assert(dlen > 0);

    uint64_t DDR_IN_ADDR = ptu.block_config[0].BASE_IN_ADDR;
    uint64_t DDR_RS_ADDR = ptu.block_config[ptu.block_config.size() - 1].BASE_RS_ADDR;

    // clear the in_data for every block
    for (int block_idx = 0; block_idx < ptu.block_config.size(); block_idx++)
    {
        ptu.block_config[block_idx].in_data.clear();
    }

    // Update the in_data for every block
    for (int block_idx = 0; block_idx < ptu.block_config.size(); block_idx++)
    {
        int pt_num = ptu.block_config[block_idx].cfg_cam_valid_num;
        std::vector<int> pt_mapping_tmp = pt_mapping[block_idx];
        ptu.block_config[block_idx].in_data.reserve(pt_num * (ptu.layer_config.cfg_inch_loop_num + 1));
        for (int j = 0; j < ptu.layer_config.cfg_inch_loop_num + 1; j++)
        {
            for (int k = 0; k < pt_num; k++)
            {
                int idx_tmp = pt_mapping_tmp[k] + j * bcfg.cfg_cam_valid_num;
                ptu.block_config[block_idx].in_data.push_back(layer1_merged_res[idx_tmp]);
            }
        }
        ptu.block_config[block_idx].BASE_IN_ADDR = DDR_IN_ADDR;
        ptu.block_config[block_idx].BASE_RS_ADDR = DDR_RS_ADDR;

#ifdef SIMULATION
        DDR_IN_ADDR += pt_num * sizeof(uint64_t) * (1 + ptu.layer_config.cfg_inch_loop_num);
        DDR_RS_ADDR += pt_num * sizeof(uint64_t) * (1 + 0);
#endif
    }

    timer_wt_init.reset();
    timer_wt_init.start();
    ptu.wt_init();
    timer_wt_init.stop();
    timer_wt_init.print("Weight init time for layer2");
    
    // Get CNN instructuin index for layer 2
    cnn_inst_index = ptu.inst_gen_block(block_pt_1024);

#ifdef SIMULATION
    ptu.compute_init();
    ptu.run();
#else
    timer_run.reset();
    timer_run.start();
    ptu.run_block(cnn_inst_index);
    timer_run.stop();
    timer_run.print("Run time for layer2");
#endif

    // Filter and merge the results
    std::vector<uint64_t> layer2_merged_res;
    merge_RESULT_feat(ptu.block_config, bcfg.cfg_cam_valid_num, 0, pt_mapping, valid_filter, layer2_merged_res);

#ifdef SIMULATION
    std::string layer2_output_file = "layer2_output.txt";
    export_merged_results(layer2_merged_res, valid_pt_num, 0, layer2_output_file);
#endif
    timer.stop();
    timer.print("Total time for SCONV");

    ///////////////////////////////////////////
    // FC layer                              //
    // ich = 8, och = 2                      //
    ///////////////////////////////////////////
    float fc_weight[16] = {-0.04190924, 0.08408088, 0.19655997, -0.33468807, 0.35655615, -0.26418057, -0.22881983, 0.03832784,
                           0.11310787, 0.27403793, 0.34556514, 0.18055771, 0.13469994, -0.10524856, -0.135478, 0.1271919};

    std::vector<int> voxel_vector;
    for (int idx = 0; idx < bcfg.cfg_cam_valid_num; idx++)
    {
        float res[2] = {0, 0};
        Data64to8 in64to8;
        in64to8.in64 = layer2_merged_res[idx];
        for (int j = 0; j < 2; j++)
        {
            for (int k = 0; k < 8; k++)
            {
                int8_t ch = in64to8.out8[k];
                res[j] += fc_weight[j * 8 + k] * ch;
            }
        }
        int classID = res[0] > res[1] ? 0 : 1;
        voxel_vector.push_back(classID);
    }

#ifdef SIMULATION
    ptu.Model_DDR_close(); // Close the simulated DDR file
#endif
}

// Deploy toolkit
    // Input: layer_configs
        // - represents a point cloud model
        // - layers can include
            // - SCONV layers (spatial convolution)
            // - FC layers (fully connected layers)
    // Output: 
    // Function: run each layer in the model
// Just like a compiler